import os
import concurrent.futures
from rembg import remove
from PIL import Image

# 📌 设置输入和输出文件夹（确保路径正确）
input_folder = r"D:/Desk/去背景"  # 存放原始图片的文件夹
output_folder = r"D:/Desk/去背景/result"  # 输出文件夹
os.makedirs(output_folder, exist_ok=True)

# 🔥 处理单张图片
def process_image(filename):
    try:
        input_path = os.path.join(input_folder, filename)
        output_path = os.path.join(output_folder, filename.replace(".jpg", ".png").replace(".jpeg", ".png"))

        # 读取图片
        input_image = Image.open(input_path)

        # 🚀 基础 `rembg` 去背景（不带优化参数）
        output_image = remove(input_image)

        # 保存透明 PNG
        output_image.save(output_path)
        print(f"✅ {filename} 处理完成！")

    except Exception as e:
        print(f"❌ 处理 {filename} 失败：{e}")

# 🔥 批量处理所有图片（多线程加速）
def batch_process():
    images = [f for f in os.listdir(input_folder) if f.lower().endswith((".jpg", ".png", ".jpeg"))]
    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
        executor.map(process_image, images)

# 🚀 运行批量处理
if __name__ == "__main__":
    batch_process()
    print("🎉 所有图片已处理完成！")
